import akshare as ak
import pandas as pd
import pymssql
from sqlalchemy import create_engine, distinct, or_, and_
from urllib.parse import quote_plus as urlquote
from datetime import datetime, date

host = '218.17.227.194'
port = '10001'
user = 'sa'
password = 'Jzt@20230510&&'
schema = 'dbo'
db = 'AlphabeeCom'


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def get_index_info():
    sql1 = "SELECT STK_CODE, RIGHT(F_INDEX_CODE, 6) FROM ETF_INDEX_INFO"
    r1 = exec_sql(sql1)
    df1 = pd.DataFrame(r1, columns=['code', 'idx'])
    df1 = df1.loc[df1['idx'] != 'unknowed']
    index_list = list(set(df1.idx.tolist()))

    sql2 = "SELECT INDEX_CODE, CONSTITUENT_CODE FROM index_constituent"
    r2 = exec_sql(sql2)
    df2 = pd.DataFrame(r2, columns=['idx', 'code'])
    stocks = list()
    indexs = list()

    for i, index_code in enumerate(index_list):
        df_p = df2.loc[(df2['idx'] == index_code)]
        if any(df_p.code.str.contains('HK')):
            continue
        if len(df_p) > 0:
            stocks.extend(df_p.code.to_list())
            indexs.append(index_code)

    stocks = list(set(stocks))
    df = df1.loc[df1['idx'].isin(indexs)]
    res = {'stocks': stocks, 'df': df}

    return res


def insert_data(stocks):
    table = "STOCK_DAY_DATA"
    sql = f"SELECT DISTINCT code FROM {table}"
    r = exec_sql(sql)
    in_codes = set([tok[0] for tok in r])
    stocks = stocks - in_codes
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    for i, code in enumerate(stocks):
        print(f'i={i}->len={len(stocks)}->code={code}')
        df = ak.stock_zh_a_hist(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
        df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                           '最低': 'low', '成交量': 'volume', '成交额': 'money', '换手率': 'hsl'}, inplace=True)

        df.insert(0, 'code', code)
        df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume', 'money', 'hsl']]
        df['datetime'] = df['datetime'].map(time_map)
        insert_db(df, table)
        # time.sleep(2)


def time_map(x):
    # print(x)
    t = datetime.strftime(x, '%Y%m%d')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


def set_date():
    st = "20230916"
    et = "20231019"
    st_d = datetime.strptime(st, '%Y%m%d').date()
    et_d = datetime.strptime(et, '%Y%m%d').date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d}

    return res


def main():
    stocks = set(get_index_info()['stocks'])
    insert_data(stocks)


if __name__ == '__main__':
    main()
